from common_import import *
from pulp import LpMaximize, LpProblem, LpVariable, lpSum, PULP_CBC_CMD

# 创建问题
model = LpProblem(name="maximize-profit-greenhouse-fields", sense=LpMaximize)

# 地块、作物、时间范围
lands = [
    "E1",
    "E2",
    "E3",
    "E4",
    "E5",
    "E6",
    "E7",
    "E8",
    "E9",
    "E10",
    "E11",
    "E12",
    "E13",
    "E14",
    "E15",
    "E16",
]
first_season_crops = [
    17,
    18,
    19,
    21,
    22,
    24,
    25,
    26,
    27,
    28,
    29,
    30,
    31,
]  # 第一季作物编号
second_season_crops = [38, 39, 40, 41]  # 第二季作物编号
years = range(2024, 2031)  # 时间范围：2024到2030
field_type = "普通大棚 "

# 面积数据
area_of_land = {land: 0.6 for land in lands}

# 每种作物的总销量（/斤）
total_sales = {
    17: 4320,
    18: 2880,
    19: 4320,
    21: 1800,
    22: 4800,
    24: 1800,
    25: 2400,
    26: 2700,
    27: 4500,
    28: 2400,
    29: 9000,
    30: 1500,
    31: 1200,
    38: 9000,
    39: 7200,
    40: 18000,
    41: 4200,
}

# 2023年的种植方案
known_planting_2023 = {
    "E1": (18, 38),
    "E2": (24, 38),
    "E3": (25, 38),
    "E4": (26, 39),
    "E5": (28, 39),
    "E6": (27, 39),
    "E7": (19, 40),
    "E8": (19, 40),
    "E9": (18, 40),
    "E10": (17, 41),
    "E11": (17, 41),
    "E12": (22, 41),
    "E13": (21, 41),
    "E14": (29, 41),
    "E15": (30, 41),
    "E16": (31, 41),
}
# 决策变量
x_first = LpVariable.dicts(
    "x_first",
    ((i, j, t) for i in lands for j in first_season_crops for t in years),
    lowBound=0,  # 最小种植面积
    cat="Continuous",
)
x_second = LpVariable.dicts(
    "x_second",
    ((i, j, t) for i in lands for j in second_season_crops for t in years),
    lowBound=0,  # 最小种植面积
    cat="Continuous",
)
z_first = LpVariable.dicts(
    "z_first",
    ((i, j, t) for i in lands for j in first_season_crops for t in years),
    cat="Binary",
)
z_second = LpVariable.dicts(
    "z_second",
    ((i, j, t) for i in lands for j in second_season_crops for t in years),
    cat="Binary",
)

# 新增变量：表示每种作物每年超出的部分
overproduction_first = LpVariable.dicts(
    "overproduction_first",
    ((j, t) for j in first_season_crops for t in years),
    lowBound=0,
    cat="Continuous",
)
overproduction_second = LpVariable.dicts(
    "overproduction_second",
    ((j, t) for j in second_season_crops for t in years),
    lowBound=0,
    cat="Continuous",
)
# 目标函数：最大化总利润，并惩罚超出部分
model += (
    lpSum(
        mapping.get_profit_idtype(j, field_type, "第一季")
        * x_first[i, j, t]
        * mapping.get_yield_idtype(j, field_type, "第一季")
        for i in lands
        for j in first_season_crops
        for t in years
    )
    + lpSum(
        mapping.get_profit_idtype(j, field_type, "第二季")
        * x_second[i, j, t]
        * mapping.get_yield_idtype(j, field_type, "第二季")
        for i in lands
        for j in second_season_crops
        for t in years
    )
    - lpSum(
        overproduction_first[j, t] * mapping.get_priceperjin_idtypeE(j)
        for j in first_season_crops
        for t in years
    )
    - lpSum(
        overproduction_second[j, t] * mapping.get_priceperjin_idtypeE(j)
        for j in second_season_crops
        for t in years
    )
)

# 约束条件1：每块地每连续三年必须种植豆类（第一季的17, 18, 19）
for i in lands:
    # 考虑2023, 2024, 2025三个连续年份
    model += (
        lpSum(
            (
                x_first[i, j, 2024]
                + x_first[i, j, 2025]
                + (area_of_land[i] if known_planting_2023[i][0] in [17, 18, 19] else 0)
            )
            for j in [17, 18, 19]
        )
        >= area_of_land[i]
    )

    # 从2024年开始，每三个连续年份至少有一个豆类作物
    for t in range(2024, 2029):
        model += (
            lpSum(x_first[i, j, t + k] for j in [17, 18, 19] for k in range(3))
            >= area_of_land[i]
        )

# 约束条件2：同一块地在连续两年内不能种植相同的作物
M = 1000
for i in lands:
    # 第一季作物的重茬约束
    for j in first_season_crops:
        for t in range(2024, 2030):
            model += x_first[i, j, t] <= M * z_first[i, j, t]
            model += z_first[i, j, t] + z_first[i, j, t + 1] <= 1

    # 第二季作物的重茬约束
    for j in second_season_crops:
        for t in range(2024, 2030):
            model += x_second[i, j, t] <= M * z_second[i, j, t]
            model += z_second[i, j, t] + z_second[i, j, t + 1] <= 1

# 2024年不能与2023年种植的作物重茬
for i, (first_crop, second_crop) in known_planting_2023.items():
    # 第一季作物不能重茬种植
    if first_crop in first_season_crops:
        model += z_first[i, first_crop, 2024] == 0
    # 第二季作物不能重茬种植
    if second_crop in second_season_crops:
        model += z_second[i, second_crop, 2024] == 0

# 约束条件3：第一季和第二季的种植面积总和不能超过地块面积
for i in lands:
    for t in years:
        # 第一季作物的种植总面积不超过地块面积
        model += lpSum(x_first[i, j, t] for j in first_season_crops) <= area_of_land[i]
        # 第二季作物的种植总面积不超过地块面积
        model += (
            lpSum(x_second[i, j, t] for j in second_season_crops) <= area_of_land[i]
        )

# 约束条件4：作物总产量限制
# 第一季作物的总产量限制
for j in first_season_crops:
    for t in years:
        total_yield_first = lpSum(
            mapping.get_yield_idtype(j, field_type, "第一季") * x_first[i, j, t]
            for i in lands
        )
        model += total_yield_first <= total_sales[j] + overproduction_first[j, t]
        model += overproduction_first[j, t] >= total_yield_first - total_sales[j]

# 第二季作物的总产量限制
for j in second_season_crops:
    for t in years:
        total_yield_second = lpSum(
            mapping.get_yield_idtype(j, field_type, "第二季") * x_second[i, j, t]
            for i in lands
        )
        model += total_yield_second <= total_sales[j] + overproduction_second[j, t]
        model += overproduction_second[j, t] >= total_yield_second - total_sales[j]


# 启用CBC求解器的详细输出
solver = PULP_CBC_CMD(msg=True, gapRel=0.001)

# 求解模型
model.solve(solver)

# 输出结果
for t in years:
    for i in lands:
        # 第一季作物输出
        for j in first_season_crops:
            if (i, j, t) in x_first and x_first[i, j, t].value() > 0:
                print(
                    f"Year {t}, Land {i}, Crop {j}: Area {x_first[i, j, t].value()} mu (First Season)"
                )

        # 第二季作物输出
        for j in second_season_crops:
            if (i, j, t) in x_second and x_second[i, j, t].value() > 0:
                print(
                    f"Year {t}, Land {i}, Crop {j}: Area {x_second[i, j, t].value()} mu (Second Season)"
                )
# 初始化一个字典来存储每年的总利润
annual_profit = {t: 0 for t in years}

# 逐年计算总利润
for t in years:
    # 计算当年的总收益
    total_revenue = sum(
        mapping.get_profit_idtype(j, field_type, "第一季")
        * x_first[i, j, t].value()
        * mapping.get_yield_idtype(j, field_type, "第一季")
        for i in lands
        for j in first_season_crops
    ) + sum(
        mapping.get_profit_idtype(j, field_type, "第二季")
        * x_second[i, j, t].value()
        * mapping.get_yield_idtype(j, field_type, "第二季")
        for i in lands
        for j in second_season_crops
    )

    # 计算当年的惩罚项（超出部分）
    total_penalty = sum(
        overproduction_first[j, t].value() * mapping.get_priceperjin_idtypeE(j)
        for j in first_season_crops
    ) + sum(
        overproduction_second[j, t].value() * mapping.get_priceperjin_idtypeE(j)
        for j in second_season_crops
    )

    # 计算每年的净利润
    annual_profit[t] = total_revenue - total_penalty

    # 输出每年的总利润
    print(f"Year {t} Total Profit: {annual_profit[t]}")

# 计算所有年份的总利润
total_profit = sum(annual_profit[t] for t in years)

# 输出所有年份的总利润
print(f"Total Profit from 2024 to 2030: {total_profit}")
